import torch
import torchvision.transforms
from PIL import Image
from model import resnet
import numpy as np
import logging
from pathlib import Path
import cv2
import glob
import os


def prediction(type_id):
    os.chdir("./")  # 日志写入地址
    logging.basicConfig(filename='example.log', level=logging.INFO,
                        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

    normalize = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    transform = torchvision.transforms.Compose([
        torchvision.transforms.RandomHorizontalFlip(),
        torchvision.transforms.RandomCrop(224, 4),
        torchvision.transforms.ToTensor(),
        normalize,
    ])

    class_dir = os.path.join("../class_model/prediction/2/")
    name_list = os.listdir(class_dir)
    all_num = len(os.listdir(class_dir))
    logging.info(f'路径{class_dir},数量{all_num}')
    print(f'路径{class_dir},数量{all_num}')
    for name_id in name_list:
        name_dir = os.path.join(class_dir, name_id)
        img = cv2.imread(name_dir, 0)
        img_PIL = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        image = transform(img_PIL)
        image = torch.unsqueeze(image, dim=0)
        # image = torch.reshape(image, (1, 3, 224, 224))
        module = torch.load(f"./model_{type_id}.th")
        model = torch.nn.DataParallel(resnet.ResNet(resnet.BasicBlock, [5, 5, 5], 3))
        model.load_state_dict(module['state_dict'])
        model.cuda()
        model.eval()
        with torch.no_grad():
            output = model(image).cpu()
            predict = torch.softmax(output, dim=1)
            predict_cla = torch.argmax(predict).numpy()
            print(f'切面{name_dir},预测{predict_cla + 1}')


# 模型推理
type_id = 36
prediction(type_id)
